import { useEffect, useRef, useState, useCallback } from 'react';
import { io, Socket } from 'socket.io-client';
import { WebSocketEvents } from '../types/api';
import { apiClient } from '../lib/api';

interface UseWebSocketOptions {
  autoConnect?: boolean;
  reconnect?: boolean;
  reconnectAttempts?: number;
  reconnectDelay?: number;
}

interface WebSocketState {
  connected: boolean;
  connecting: boolean;
  error: string | null;
  socket: Socket | null;
}

export function useWebSocket(options: UseWebSocketOptions = {}) {
  const {
    autoConnect = true,
    reconnect = true,
    reconnectAttempts = 5,
    reconnectDelay = 1000,
  } = options;

  const [state, setState] = useState<WebSocketState>({
    connected: false,
    connecting: false,
    error: null,
    socket: null,
  });

  const socketRef = useRef<Socket | null>(null);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const reconnectCountRef = useRef(0);

  const connect = useCallback(() => {
    if (socketRef.current?.connected) {
      return;
    }

    setState(prev => ({ ...prev, connecting: true, error: null }));

    const socket = io(apiClient.getWebSocketUrl(), {
      transports: ['websocket'],
      upgrade: false,
      rememberUpgrade: false,
    });

    socket.on('connect', () => {
      console.log('WebSocket connected');
      reconnectCountRef.current = 0;
      setState(prev => ({
        ...prev,
        connected: true,
        connecting: false,
        error: null,
        socket,
      }));
    });

    socket.on('disconnect', (reason) => {
      console.log('WebSocket disconnected:', reason);
      setState(prev => ({
        ...prev,
        connected: false,
        connecting: false,
        socket: null,
      }));

      // Auto-reconnect logic
      if (reconnect && reconnectCountRef.current < reconnectAttempts) {
        reconnectCountRef.current++;
        console.log(`Attempting to reconnect (${reconnectCountRef.current}/${reconnectAttempts})...`);
        
        reconnectTimeoutRef.current = setTimeout(() => {
          connect();
        }, reconnectDelay * reconnectCountRef.current);
      }
    });

    socket.on('connect_error', (error) => {
      console.error('WebSocket connection error:', error);
      setState(prev => ({
        ...prev,
        connected: false,
        connecting: false,
        error: error.message,
        socket: null,
      }));
    });

    // Handle specific events
    socket.on('connected', (data: WebSocketEvents['connected']) => {
      console.log('Server connection confirmed:', data);
    });

    socket.on('error', (data: WebSocketEvents['error']) => {
      console.error('Server error:', data.message);
      setState(prev => ({ ...prev, error: data.message }));
    });

    socketRef.current = socket;
  }, [reconnect, reconnectAttempts, reconnectDelay]);

  const disconnect = useCallback(() => {
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
      reconnectTimeoutRef.current = null;
    }

    if (socketRef.current) {
      socketRef.current.disconnect();
      socketRef.current = null;
    }

    setState({
      connected: false,
      connecting: false,
      error: null,
      socket: null,
    });
  }, []);

  // Event subscription helpers
  const subscribe = useCallback(<T extends keyof WebSocketEvents>(
    event: T,
    handler: (data: WebSocketEvents[T]) => void
  ) => {
    if (socketRef.current) {
      socketRef.current.on(event as string, handler as any);
      return () => {
        if (socketRef.current) {
          socketRef.current.off(event as string, handler as any);
        }
      };
    }
    return () => {};
  }, []);

  const emit = useCallback(<T extends keyof WebSocketEvents>(
    event: T,
    data: WebSocketEvents[T]
  ) => {
    if (socketRef.current?.connected) {
      socketRef.current.emit(event, data);
    } else {
      console.warn(`Cannot emit ${event}: WebSocket not connected`);
    }
  }, []);

  // Authentication
  const authenticate = useCallback((publicKey: string) => {
    emit('authenticate', { publicKey });
  }, [emit]);

  // Token subscriptions
  const subscribeToToken = useCallback((mint: string) => {
    emit('subscribe_token', { mint });
  }, [emit]);

  const unsubscribeFromToken = useCallback((mint: string) => {
    emit('unsubscribe_token', { mint });
  }, [emit]);

  // Address subscriptions
  const subscribeToAddress = useCallback((address: string) => {
    emit('subscribe_address', { address });
  }, [emit]);

  // Auto-connect on mount
  useEffect(() => {
    if (autoConnect) {
      connect();
    }

    return () => {
      disconnect();
    };
  }, [autoConnect, connect, disconnect]);

  return {
    // State
    connected: state.connected,
    connecting: state.connecting,
    error: state.error,

    // Connection control
    connect,
    disconnect,

    // Event handling
    subscribe,
    emit,

    // Convenience methods
    authenticate,
    subscribeToToken,
    unsubscribeFromToken,
    subscribeToAddress,
  };
}

// Hook for specific token data
export function useTokenWebSocket(mint?: string) {
  const ws = useWebSocket();
  const [tokenData, setTokenData] = useState<any>(null);
  const [priceUpdates, setPriceUpdates] = useState<any[]>([]);
  const [transactions, setTransactions] = useState<any[]>([]);

  // Subscribe to token when mint is available
  useEffect(() => {
    if (ws.connected && mint) {
      ws.subscribeToToken(mint);

      const unsubscribePriceUpdate = ws.subscribe('token_price_update', (data) => {
        if (data.mint === mint) {
          setPriceUpdates(prev => [...prev.slice(-99), data]);
        }
      });

      const unsubscribeTransaction = ws.subscribe('token_transaction', (data) => {
        if (data.transaction.tokenMint === mint) {
          setTransactions(prev => [data.transaction, ...prev.slice(0, 49)]);
        }
      });

      return () => {
        ws.unsubscribeFromToken(mint);
        unsubscribePriceUpdate();
        unsubscribeTransaction();
      };
    }
  }, [ws.connected, mint, ws]);

  return {
    ...ws,
    tokenData,
    priceUpdates,
    transactions,
  };
}
